import json
import os
import shutil
import subprocess
import sys

import requests


def download_json(url):
    # 下载JSON文件
    response = requests.get(url)
    response.raise_for_status()  # 检查请求是否成功
    return response.json()


def download_and_modify_json(url, local_filename, modifications):
    if os.path.exists(local_filename):
        data = json.load(open(local_filename))
        config_version = data.get('config_version', '0.0.0')
        if config_version < '1.2.0':
            data = download_json(url)
    else:
        data = download_json(url)

    # 修改内容
    for key, value in modifications.items():
        data[key] = value

    # 保存修改后的内容
    with open(local_filename, 'w', encoding='utf-8') as f:
        json.dump(data, f, ensure_ascii=False, indent=4)


def git_clone_model(repo_id, output_dir, patterns=None, use_mirror=True):
    """
    使用git clone下载模型
    """
    print(f"正在使用git clone下载模型: {repo_id}")
    
    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)
    
    # 设置环境变量，跳过LFS文件的自动下载
    env = os.environ.copy()
    env['GIT_LFS_SKIP_SMUDGE'] = '1'
    
    # 构建git clone命令，使用镜像站点
    if use_mirror:
        repo_url = f"https://hf-mirror.com/{repo_id}"
        print(f"使用镜像站点: https://hf-mirror.com")
    else:
        repo_url = f"https://huggingface.co/{repo_id}"
    
    cmd = ["git", "clone", repo_url, output_dir]
    
    try:
        # 执行git clone命令
        subprocess.run(cmd, env=env, check=True)
        print(f"成功克隆仓库: {repo_id}")
        
        # 进入目录并安装git lfs
        cwd = os.getcwd()
        os.chdir(output_dir)
        
        # 安装git lfs
        subprocess.run(["git", "lfs", "install"], check=True)
        
        # 如果指定了patterns，只下载特定文件
        if patterns:
            for pattern in patterns:
                # 构建git lfs pull命令，只拉取特定文件
                lfs_cmd = ["git", "lfs", "pull", "--include", pattern]
                subprocess.run(lfs_cmd, check=True)
                print(f"已下载文件: {pattern}")
        else:
            # 下载所有LFS文件
            subprocess.run(["git", "lfs", "pull"], check=True)
            print("已下载所有LFS文件")
        
        # 返回原目录
        os.chdir(cwd)
        
        return output_dir
    except subprocess.CalledProcessError as e:
        print(f"下载模型时出错: {e}")
        if use_mirror:
            print("镜像站点访问失败，尝试使用原始URL...")
            return git_clone_model(repo_id, output_dir, patterns, use_mirror=False)
        else:
            print("请检查网络连接或尝试使用代理。")
            return None


if __name__ == '__main__':
    # 设置模型下载目录
    base_dir = os.path.expanduser("~/models")
    os.makedirs(base_dir, exist_ok=True)
    
    # PDF-Extract-Kit模型
    pdf_extract_kit_dir = os.path.join(base_dir, "PDF-Extract-Kit-1.0")
    mineru_patterns = [
        "models/Layout/YOLO/*",
        "models/MFD/YOLO/*",
        "models/MFR/unimernet_hf_small_2503/*",
        "models/OCR/paddleocr_torch/*",
    ]
    
    # 使用git clone下载PDF-Extract-Kit模型
    try:
        model_dir = git_clone_model("opendatalab/PDF-Extract-Kit-1.0", pdf_extract_kit_dir, mineru_patterns)
        
        if not model_dir:
            sys.exit(1)
        
        # layoutreader模型
        layoutreader_dir = os.path.join(base_dir, "layoutreader")
        layoutreader_patterns = [
            "*.json",
            "*.safetensors",
        ]
        
        # 使用git clone下载layoutreader模型
        layoutreader_model_dir = git_clone_model("hantian/layoutreader", layoutreader_dir, layoutreader_patterns)
        
        if not layoutreader_model_dir:
            sys.exit(1)
        
        # 设置模型目录路径
        model_dir = os.path.join(model_dir, "models")
        print(f'model_dir is: {model_dir}')
        print(f'layoutreader_model_dir is: {layoutreader_model_dir}')
        
        # 下载并修改配置文件
        json_url = 'https://github.com/opendatalab/MinerU/raw/master/magic-pdf.template.json'
        config_file_name = 'magic-pdf.json'
        home_dir = os.path.expanduser('~')
        config_file = os.path.join(home_dir, config_file_name)
        
        json_mods = {
            'models-dir': model_dir,
            'layoutreader-model-dir': layoutreader_model_dir,
        }
        
        download_and_modify_json(json_url, config_file, json_mods)
        print(f'配置文件已成功配置，路径为: {config_file}')
        
    except Exception as e:
        print(f'执行过程中出错: {e}')
        print('请检查网络连接或尝试使用代理。')
